#include <stdio.h>
#include <string.h>
#include <gmssl/sm4.h>
#include <gmssl/rand.h>
#include <stdint.h> // 包含 uint8_t

// 定义SM4密钥和IV的固定大小和最大缓冲区
#define SM4_KEY_SIZE 16
#define SM4_BLOCK_SIZE 16
#define MAX_BUFFER_SIZE 128

// (复用 print_hex 函数)
void print_hex(const char* label, const unsigned char* data, int len) {
    printf("%s (%d bytes):\n", label, len);
    for (int i = 0; i < len; ++i) {
        printf("%02x", data[i]);
    }
    printf("\n");
}

int main() {
    // 原始密钥数据，用于初始化加密/解密上下文
    uint8_t raw_key[SM4_KEY_SIZE]; 
    // 初始向量
    uint8_t iv[SM4_BLOCK_SIZE];
    
    // 明文数据
    // 移除了原始代码中关于对齐的误导注释，GMSSL上下文接口会自动处理 PKCS#7 填充
    const uint8_t *plaintext = (const uint8_t *)"This is a test message for SM4 with GmSSL CBC Mode.";
    size_t plaintext_len = strlen((char*)plaintext); 

    // 密文和解密后文本的缓冲区，需要足够大以容纳明文+填充
    uint8_t ciphertext[MAX_BUFFER_SIZE];
    uint8_t decryptedtext[MAX_BUFFER_SIZE];
    
    size_t ciphertext_len = 0;
    size_t decrypted_len = 0;
    
    // -----------------------------------------------------
    // 1. 生成密钥和IV
    // -----------------------------------------------------
    if (rand_bytes(raw_key, SM4_KEY_SIZE) != 1) {
        fprintf(stderr, "Error generating random key.\n");
        return 1;
    }
    if (rand_bytes(iv, SM4_BLOCK_SIZE) != 1) {
        fprintf(stderr, "Error generating random IV.\n");
        return 1;
    }

    print_hex("SM4 Key", raw_key, SM4_KEY_SIZE);
    print_hex("SM4 IV (Encrypt)", iv, SM4_BLOCK_SIZE);

    // -----------------------------------------------------
    // 2. 加密 (CBC) - 使用 SM4_CBC_CTX 上下文
    // -----------------------------------------------------
    SM4_CBC_CTX ctx;
    size_t update_len = 0;
    size_t final_len = 0;
    
    // 初始化加密上下文
    sm4_cbc_encrypt_init(&ctx, raw_key, iv);

    // 更新加密数据 (处理大部分数据块)
    sm4_cbc_encrypt_update(&ctx, plaintext, plaintext_len, ciphertext, &update_len);

    // 完成加密 (处理最后的填充块，并计算最终长度)
    sm4_cbc_encrypt_finish(&ctx, ciphertext + update_len, &final_len);

    ciphertext_len = update_len + final_len;

    print_hex("Ciphertext", ciphertext, (int)ciphertext_len);

    // -----------------------------------------------------
    // 3. 解密 (CBC) - 使用 SM4_CBC_CTX 上下文
    // -----------------------------------------------------
    SM4_CBC_CTX dec_ctx;
    size_t dec_update_len = 0;
    size_t dec_final_len = 0;

    // 解密时必须使用原始 IV
    uint8_t iv_dec[SM4_BLOCK_SIZE];
    memcpy(iv_dec, iv, SM4_BLOCK_SIZE); 
    
    // 初始化解密上下文
    sm4_cbc_decrypt_init(&dec_ctx, raw_key, iv_dec);

    // 更新解密数据 
    sm4_cbc_decrypt_update(&dec_ctx, ciphertext, ciphertext_len, decryptedtext, &dec_update_len);

    // 完成解密 (移除填充，并返回最终有效长度)
    sm4_cbc_decrypt_finish(&dec_ctx, decryptedtext + dec_update_len, &dec_final_len);
    
    decrypted_len = dec_update_len + dec_final_len;
    
    // -----------------------------------------------------
    // 4. 验证结果
    // -----------------------------------------------------
    
    // 在解密后的文本末尾添加 null 终止符，以便打印
    if (decrypted_len < MAX_BUFFER_SIZE) {
        decryptedtext[decrypted_len] = '\0';
    }

    printf("\nOriginal length: %zu bytes\n", plaintext_len);
    printf("Decrypted length: %zu bytes\n", decrypted_len);

    if (decrypted_len == plaintext_len && memcmp(plaintext, decryptedtext, plaintext_len) == 0) {
        printf("Decrypted text: %s\n", decryptedtext);
        printf("GmSSL SM4 Enc/Dec Success!\n");
    } else {
        printf("GmSSL SM4 Enc/Dec Failed!\n");
        print_hex("Original Plaintext", plaintext, (int)plaintext_len);
        print_hex("Decrypted Plaintext", decryptedtext, (int)decrypted_len);
    }

    return 0;
}

